package main

import (
	"net/url"
	"sync"
	"sync/atomic"
)

// LoadBalancer 负载均衡器接口
type LoadBalancer interface {
	// Select 选择一个上游服务器
	Select() (*url.URL, error)
}

// RoundRobinLB 轮询负载均衡器
type RoundRobinLB struct {
	servers []*url.URL
	current uint64 // 当前索引（原子操作）
}

// NewRoundRobinLB 创建轮询负载均衡器
func NewRoundRobinLB(servers []UpstreamServer) (*RoundRobinLB, error) {
	urls := make([]*url.URL, 0, len(servers))
	for _, srv := range servers {
		u, err := url.Parse(srv.Target)
		if err != nil {
			continue // 跳过无效的 URL
		}
		urls = append(urls, u)
	}
	if len(urls) == 0 {
		return nil, nil
	}
	return &RoundRobinLB{servers: urls}, nil
}

func (lb *RoundRobinLB) Select() (*url.URL, error) {
	if len(lb.servers) == 0 {
		return nil, nil
	}
	idx := atomic.AddUint64(&lb.current, 1) - 1
	return lb.servers[idx%uint64(len(lb.servers))], nil
}

// WeightedLB 加权负载均衡器
type WeightedLB struct {
	servers []weightedServer
	total   int
	mu      sync.Mutex
}

type weightedServer struct {
	url     *url.URL
	weight  int
	current int // 当前权重
}

// NewWeightedLB 创建加权负载均衡器
func NewWeightedLB(servers []UpstreamServer) (*WeightedLB, error) {
	weighted := make([]weightedServer, 0, len(servers))
	total := 0
	for _, srv := range servers {
		u, err := url.Parse(srv.Target)
		if err != nil {
			continue // 跳过无效的 URL
		}
		weight := srv.Weight
		if weight <= 0 {
			weight = 1
		}
		weighted = append(weighted, weightedServer{
			url:     u,
			weight:  weight,
			current: weight,
		})
		total += weight
	}
	if len(weighted) == 0 {
		return nil, nil
	}
	return &WeightedLB{
		servers: weighted,
		total:   total,
	}, nil
}

func (lb *WeightedLB) Select() (*url.URL, error) {
	if len(lb.servers) == 0 {
		return nil, nil
	}

	lb.mu.Lock()
	defer lb.mu.Unlock()

	// 加权轮询算法（WRR）
	// 每次选择当前权重最高的服务器，然后减少其当前权重
	maxIdx := 0
	maxCurrent := lb.servers[0].current

	for i := 1; i < len(lb.servers); i++ {
		if lb.servers[i].current > maxCurrent {
			maxCurrent = lb.servers[i].current
			maxIdx = i
		}
	}

	selected := lb.servers[maxIdx]

	// 减少选中服务器的当前权重
	selected.current -= lb.total

	// 如果所有服务器的当前权重都降到0或以下，重置为原始权重
	allZero := true
	for i := range lb.servers {
		if lb.servers[i].current > 0 {
			allZero = false
			break
		}
	}

	if allZero {
		for i := range lb.servers {
			lb.servers[i].current = lb.servers[i].weight
		}
	}

	return selected.url, nil
}

// CreateLoadBalancer 根据算法类型创建负载均衡器
func CreateLoadBalancer(up Upstream) (LoadBalancer, error) {
	if len(up.Servers) == 0 {
		return nil, nil
	}

	algorithm := up.Algorithm
	if algorithm == "" {
		algorithm = "round-robin"
	}

	switch algorithm {
	case "weight", "weighted":
		return NewWeightedLB(up.Servers)
	case "round-robin", "roundrobin":
		return NewRoundRobinLB(up.Servers)
	default:
		return NewRoundRobinLB(up.Servers)
	}
}
